/*
 * Copyright (c) Kumo Inc. and affiliates.
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include <melon/singleton.h>

#ifndef _WIN32
#include <dlfcn.h>
#include <signal.h>
#include <time.h>
#endif

#include <atomic>
#include <chrono>
#include <cstdio>
#include <cstdlib>
#include <string>
#include <fmt/chrono.h>
#include <fmt/format.h>

#include <melon/demangle.h>
#include <melon/scope_guard.h>
#include <melon/experimental/symbolizer/symbolizer.h>
#include <melon/lang/safe_assert.h>
#include <melon/portability/config.h>
#include <melon/portability/fmt_compile.h>
// Before registrationComplete() we cannot assume that glog has been
// initialized, so we need to use RAW_LOG for any message that may be logged
// before that.
#include <turbo/base/internal/raw_logging.h>

#if !defined(_WIN32) && !defined(__APPLE__) && !defined(__ANDROID__)
#define MELON_SINGLETON_HAVE_DLSYM 1
#endif

namespace melon {
#if MELON_SINGLETON_HAVE_DLSYM
    namespace detail {
        static void singleton_hs_init_weak(int *argc, char **argv[])
        __attribute__((__weakref__("hs_init")));
    } // namespace detail
#endif

    SingletonVault::Type SingletonVault::defaultVaultType() {
#if MELON_SINGLETON_HAVE_DLSYM
        bool isPython = dlsym(RTLD_DEFAULT, "Py_Main");
        bool isHaskell =
                detail::singleton_hs_init_weak || dlsym(RTLD_DEFAULT, "hs_init");
        bool isJVM = dlsym(RTLD_DEFAULT, "JNI_GetCreatedJavaVMs");
        bool isD = dlsym(RTLD_DEFAULT, "_d_run_main");

        return isPython || isHaskell || isJVM || isD ? Type::Relaxed : Type::Strict;
#else
  return Type::Relaxed;
#endif
    }

    namespace detail {
        std::string TypeDescriptor::name() const {
            auto ret = demangle(ti_.name());
            if (tag_ti_ != std::type_index(typeid(DefaultTag))) {
                ret += "/";
                ret += demangle(tag_ti_.name());
            }
            return ret.toStdString();
        }

        [[noreturn]] void singletonWarnDoubleRegistrationAndAbort(
            const TypeDescriptor &type) {
            MELON_SAFE_FATAL( // May happen before registrationComplete().
                fmt::format(
                    fmt::runtime( //
                        "Double registration of singletons of the same "
                        "underlying type; check for multiple definitions "
                        "of type melon::Singleton<{}>"),
                    type.name())
                .c_str());
        }

        [[noreturn]] void singletonWarnLeakyDoubleRegistrationAndAbort(
            const TypeDescriptor &type) {
            MELON_SAFE_FATAL( // May happen before registrationComplete().
                fmt::format(
                    fmt::runtime( //
                        "Double registration of singletons of the same "
                        "underlying type; check for multiple definitions "
                        "of type melon::LeakySingleton<{}>"),
                    type.name())
                .c_str());
        }

        [[noreturn]] void singletonWarnLeakyInstantiatingNotRegisteredAndAbort(
            const TypeDescriptor &type) {
            MELON_SAFE_FATAL( // May happen before registrationComplete().
                "Creating instance for unregistered singleton: ",
                type.name().c_str());
        }

        [[noreturn]] void singletonWarnRegisterMockEarlyAndAbort(
            const TypeDescriptor &type) {
            MELON_SAFE_FATAL( // May happen before registrationComplete().
                "Registering mock before singleton was registered: ",
                type.name().c_str());
        }

        void singletonWarnDestroyInstanceLeak(
            const TypeDescriptor &type, const void *ptr) {
            KLOG(ERROR) << "Singleton of type " << type.name() << " has a "
                    << "living reference at destroyInstances time; beware! Raw "
                    << "pointer is " << ptr << ". It is very likely "
                    << "that some other singleton is holding a shared_ptr to it. "
                    << "This singleton will be leaked (even if a shared_ptr to it "
                    << "is eventually released)."
                    << "Make sure dependencies between these singletons are "
                    << "properly defined.";
        }

        [[noreturn]] void singletonWarnCreateCircularDependencyAndAbort(
            const TypeDescriptor &type) {
            MELON_SAFE_FATAL("circular singleton dependency: ", type.name().c_str());
        }

        [[noreturn]] void singletonWarnCreateUnregisteredAndAbort(
            const TypeDescriptor &type) {
            MELON_SAFE_FATAL( // May happen before registrationComplete().
                "Creating instance for unregistered singleton: ",
                type.name().c_str());
        }

        [[noreturn]] void singletonWarnCreateBeforeRegistrationCompleteAndAbort(
            const TypeDescriptor &type) {
            MELON_SAFE_FATAL( // May happen before registrationComplete().
                fmt::format(
                    fmt::runtime( //
                        "Singleton {} requested before "
                        "registrationComplete() call.\n"
                        "This usually means that either main() never called "
                        "melon::init, or singleton was requested before main() "
                        "(which is not allowed)"),
                    type.name())
                .c_str());
        }

        void singletonPrintDestructionStackTrace(const TypeDescriptor &type) {
            auto trace = symbolizer::getStackTraceStr();
            KLOG(ERROR) << "Singleton " << type.name() << " was released.\n"
                    << "Stacktrace:\n"
                    << (!trace.empty() ? trace : "(not available)");
        }

        [[noreturn]] void singletonThrowNullCreator(const std::type_info &type) {
            auto const msg = fmt::format(
                MELON_FMT_COMPILE(
                    "nullptr_t should be passed if you want {} to be default constructed"),
                melon::StringPiece(demangle(type)));
            throw std::logic_error(msg);
        }

        [[noreturn]] void singletonThrowGetInvokedAfterDestruction(
            const TypeDescriptor &type) {
            throw std::runtime_error(
                "Raw pointer to a singleton requested after its destruction."
                " Singleton type is: " +
                type.name());
        }
    } // namespace detail

    namespace {
        struct FatalHelper {
            ~FatalHelper() {
                if (!leakedSingletons_.empty()) {
                    std::string leakedTypes;
                    for (const auto &singleton: leakedSingletons_) {
                        leakedTypes += "\t" + singleton.name() + "\n";
                    }
                    KLOG(DFATAL) << "Singletons of the following types had living references "
                            << "after destroyInstances was finished:\n"
                            << leakedTypes
                            << "beware! It is very likely that those singleton instances "
                            << "are leaked.";
                }
            }

            std::vector<detail::TypeDescriptor> leakedSingletons_;
        };

#if defined(__APPLE__) || defined(_MSC_VER)
// OS X doesn't support constructor priorities.
FatalHelper fatalHelper;
#else
        FatalHelper __attribute__((__init_priority__(101))) fatalHelper;
#endif
    } // namespace

    SingletonVault::SingletonVault(Type type) noexcept : type_(type) {
        AtFork::registerHandler(
            this,
            /*prepare*/
            [this]() {
                auto singletons = singletons_.rlock();
                auto creationOrder = creationOrder_.rlock();

                KCHECK_GE(singletons->size(), creationOrder->size());

                for (const auto &singletonType: *creationOrder) {
                    liveSingletonsPreFork_.insert(singletons->at(singletonType));
                }

                return true;
            },
            /*parent*/ [this]() { liveSingletonsPreFork_.clear(); },
            /*child*/
            [this]() {
                for (auto singleton: liveSingletonsPreFork_) {
                    singleton->inChildAfterFork();
                }
                liveSingletonsPreFork_.clear();
            });
    }

    SingletonVault::~SingletonVault() {
        AtFork::unregisterHandler(this);
        destroyInstances();
    }

    void SingletonVault::registerSingleton(detail::SingletonHolderBase *entry) {
        auto state = state_.rlock();
        state->check(detail::SingletonVaultState::Type::Running);

        if (MELON_UNLIKELY(state->registrationComplete) &&
            type_.load(std::memory_order_relaxed) == Type::Strict) {
            KLOG(ERROR) << "Registering singleton after registrationComplete().";
        }

        auto singletons = singletons_.wlock();
        CHECK_THROW(
            singletons->emplace(entry->type(), entry).second, std::logic_error);
    }

    void SingletonVault::addEagerInitSingleton(detail::SingletonHolderBase *entry) {
        auto state = state_.rlock();
        state->check(detail::SingletonVaultState::Type::Running);

        if (MELON_UNLIKELY(state->registrationComplete) &&
            type_.load(std::memory_order_relaxed) == Type::Strict) {
            KLOG(ERROR) << "Registering for eager-load after registrationComplete().";
        }

        CHECK_THROW(singletons_.rlock()->count(entry->type()), std::logic_error);

        auto eagerInitSingletons = eagerInitSingletons_.wlock();
        eagerInitSingletons->insert(entry);
    }

    void SingletonVault::addEagerInitOnReenableSingleton(
        detail::SingletonHolderBase *entry) {
        auto state = state_.rlock();
        state->check(detail::SingletonVaultState::Type::Running);

        if (MELON_UNLIKELY(state->registrationComplete) &&
            type_.load(std::memory_order_relaxed) == Type::Strict) {
            KLOG(ERROR)
                    << "Registering for eager-load on re-enable after registrationComplete().";
        }

        CHECK_THROW(singletons_.rlock()->count(entry->type()), std::logic_error);

        auto eagerInitOnReenableSingletons = eagerInitOnReenableSingletons_.wlock();
        eagerInitOnReenableSingletons->insert(entry);
    }

    void SingletonVault::registrationComplete() {
        scheduleDestroyInstances();

        auto state = state_.wlock();
        state->check(detail::SingletonVaultState::Type::Running);

        if (state->registrationComplete) {
            return;
        }

        auto singletons = singletons_.rlock();
        if (type_.load(std::memory_order_relaxed) == Type::Strict) {
            for (const auto &p: *singletons) {
                if (p.second->hasLiveInstance()) {
                    throw std::runtime_error(
                        "Singleton " + p.first.name() +
                        " created before registration was complete.");
                }
            }
        }

        state->registrationComplete = true;
    }

    void SingletonVault::doEagerInit() { {
            auto state = state_.rlock();
            state->check(detail::SingletonVaultState::Type::Running);
            if (MELON_UNLIKELY(!state->registrationComplete)) {
                throw std::logic_error("registrationComplete() not yet called");
            }
        }

        auto eagerInitSingletons = eagerInitSingletons_.rlock();
        for (auto *single: *eagerInitSingletons) {
            single->createInstance();
        }
    }

    void SingletonVault::doEagerInitVia(Executor &exe, melon::Baton<> *done) { {
            auto state = state_.rlock();
            state->check(detail::SingletonVaultState::Type::Running);
            if (MELON_UNLIKELY(!state->registrationComplete)) {
                throw std::logic_error("registrationComplete() not yet called");
            }
        }

        auto eagerInitSingletons = eagerInitSingletons_.rlock();
        auto countdown =
                std::make_shared<std::atomic<size_t> >(eagerInitSingletons->size());
        for (auto *single: *eagerInitSingletons) {
            // countdown is retained by shared_ptr, and will be alive until last lambda
            // is done.  notifyBaton is provided by the caller, and expected to remain
            // present (if it's non-nullptr).  singletonSet can go out of scope but
            // its values, which are SingletonHolderBase pointers, are alive as long as
            // SingletonVault is not being destroyed.
            exe.add([=] {
                // decrement counter and notify if requested, whether initialization
                // was successful, was skipped (already initialized), or exception thrown.
                SCOPE_EXIT {
                    if (--(*countdown) == 0) {
                        if (done != nullptr) {
                            done->post();
                        }
                    }
                };
                // if initialization is in progress in another thread, don't try to init
                // here.  Otherwise the current thread will block on 'createInstance'.
                if (!single->creationStarted()) {
                    single->createInstance();
                }
            });
        }
    }

    void SingletonVault::destroyInstances() {
        cancellationSource_.wlock()->requestCancellation();

        auto stateW = state_.wlock();
        if (stateW->state == detail::SingletonVaultState::Type::Quiescing) {
            return;
        }
        stateW->state = detail::SingletonVaultState::Type::Quiescing;

        auto stateR = stateW.moveFromWriteToRead(); {
            auto singletons = singletons_.rlock();
            auto creationOrder = creationOrder_.rlock();

            KCHECK_GE(singletons->size(), creationOrder->size());

            // Release all ReadMostlyMainPtrs at once
            {
                ReadMostlyMainPtrDeleter<> deleter;
                for (auto &singleton_type: *creationOrder) {
                    singletons->at(singleton_type)->preDestroyInstance(deleter);
                }
            }

            for (auto type_iter = creationOrder->rbegin();
                 type_iter != creationOrder->rend();
                 ++type_iter) {
                singletons->at(*type_iter)->destroyInstance();
            }

            for (auto &singleton_type: *creationOrder) {
                auto instance = singletons->at(singleton_type);
                if (!instance->hasLiveInstance()) {
                    continue;
                }

                fatalHelper.leakedSingletons_.push_back(instance->type());
            }
        } {
            auto creationOrder = creationOrder_.wlock();
            creationOrder->clear();
        }
    }

    void SingletonVault::reenableInstances() {
        KCHECK(!shutdownTimerStarted_.load(std::memory_order_relaxed))
      << "reenableInstances() called after destroyInstancesFinal()"; {
            auto state = state_.wlock();

            state->check(detail::SingletonVaultState::Type::Quiescing);

            state->state = detail::SingletonVaultState::Type::Running;
        }

        // reset the cancellation source
        cancellationSource_.withWLock([&](auto &cancellationSource) {
            cancellationSource = melon::CancellationSource{};
        });

        auto eagerInitOnReenableSingletons = eagerInitOnReenableSingletons_.copy();
        auto instantiatedAtLeastOnce = instantiatedAtLeastOnce_.copy();
        for (auto *single: eagerInitOnReenableSingletons) {
            if (!instantiatedAtLeastOnce.count(single->type())) {
                continue;
            }
            single->createInstance();
        }
    }

    void SingletonVault::scheduleDestroyInstances() {
        // Add a dependency on melon::ThreadLocal to make sure all its static
        // singletons are initalized first.
        threadlocal_detail::StaticMeta<void, void>::instance();
#if !defined(MELON_SINGLETON_SKIP_SCHEDULE_ATEXIT) || \
    !MELON_SINGLETON_SKIP_SCHEDULE_ATEXIT
        std::atexit([] { SingletonVault::singleton()->destroyInstancesFinal(); });
#endif
    }

    void SingletonVault::destroyInstancesFinal() {
        startShutdownTimer();
        destroyInstances();
    }

    void SingletonVault::addToShutdownLog(std::string message) {
        std::chrono::time_point<std::chrono::system_clock> now =
                std::chrono::system_clock::now();
        std::chrono::milliseconds millis =
                std::chrono::duration_cast<std::chrono::milliseconds>(
                    now.time_since_epoch());
        shutdownLog_.wlock()->push_back(fmt::format("{:%T} {}", millis, message));
    }

#if MELON_HAVE_LIBRT
namespace {
[[noreturn]] void fireShutdownSignalHelper(sigval_t sigval) {
  static_cast<SingletonVault*>(sigval.sival_ptr)->fireShutdownTimer();
}
} // namespace
#endif

    void SingletonVault::startShutdownTimer() {
#if MELON_HAVE_LIBRT
  if (shutdownTimerStarted_.exchange(true)) {
    return;
  }

  if (!shutdownTimeout_.count()) {
    return;
  }

  struct sigevent sig;
  sig.sigev_notify = SIGEV_THREAD;
  sig.sigev_notify_function = fireShutdownSignalHelper;
  sig.sigev_value.sival_ptr = this;
  sig.sigev_notify_attributes = nullptr;
  timer_t timerId;
  PKCHECK(timer_create(CLOCK_MONOTONIC, &sig, &timerId) == 0);

  struct itimerspec newValue, oldValue;
  newValue.it_value.tv_sec =
      std::chrono::milliseconds(shutdownTimeout_).count() / 1000;
  newValue.it_value.tv_nsec =
      std::chrono::milliseconds(shutdownTimeout_).count() % 1000 * 1000000;
  newValue.it_interval.tv_sec = 0;
  newValue.it_interval.tv_nsec = 0;
  PKCHECK(timer_settime(timerId, 0, &newValue, &oldValue) == 0);
#endif
    }

    [[noreturn]] void SingletonVault::fireShutdownTimer() {
        std::string shutdownLog;
        for (auto &logMessage: shutdownLog_.copy()) {
            shutdownLog += logMessage + "\n";
        }

        auto msg = melon::to<std::string>(
            "Failed to complete shutdown within ",
            std::chrono::milliseconds(shutdownTimeout_).count(),
            "ms. Shutdown log:\n",
            shutdownLog);
        melon::terminate_with<std::runtime_error>(msg);
    }
} // namespace melon
